package org.baicaixiaozhan.utils;

import com.google.common.collect.Multimap;
import com.google.common.collect.MultimapBuilder;
import com.google.common.io.ByteSource;

import javax.crypto.*;
import javax.crypto.spec.GCMParameterSpec;
import javax.crypto.spec.PBEKeySpec;
import javax.crypto.spec.SecretKeySpec;
import java.io.*;
import java.nio.ByteBuffer;
import java.nio.charset.StandardCharsets;
import java.security.InvalidAlgorithmParameterException;
import java.security.InvalidKeyException;
import java.security.NoSuchAlgorithmException;
import java.security.SecureRandom;
import java.security.spec.InvalidKeySpecException;
import java.security.spec.KeySpec;
import java.util.*;

import static java.util.stream.Collectors.joining;
import static java.util.stream.Collectors.toList;

/**
 * DESC:
 *
 * @author : baicaixiaozhan
 * @since : 2020-12-20
 */
public final class AES256GCMUtils {

    private static final int AES_KEY_SIZE = 256;
    private static final int GCM_IV_LENGTH = 12;
    private static final int GCM_TAG_LENGTH = 16;
    private static final String KEY = "key";

    private AES256GCMUtils() {
        throw new AssertionError();
    }

    /**
     * This method will encrypt the given data
     *
     * @param key  : the password that will be used to encrypt the data
     * @param data : the data that will be encrypted
     * @return Encrypted data in a byte array
     */
    public static byte[] encryptData(String key, byte[] data) throws NoSuchPaddingException,
            NoSuchAlgorithmException,
            InvalidAlgorithmParameterException,
            InvalidKeyException,
            BadPaddingException,
            IllegalBlockSizeException, InvalidKeySpecException {

        //Prepare the nonce
        SecureRandom secureRandom = new SecureRandom();

        //Nonce should be 12 bytes
        byte[] iv = new byte[GCM_IV_LENGTH];
        secureRandom.nextBytes(iv);

        //Prepare your key/password
        SecretKey secretKey = generateSecretKey(key, iv);


        Cipher cipher = Cipher.getInstance("AES/GCM/NoPadding");
        GCMParameterSpec parameterSpec = new GCMParameterSpec(GCM_TAG_LENGTH * 8, iv);

        //Encryption mode on!
        cipher.init(Cipher.ENCRYPT_MODE, secretKey, parameterSpec);

        //Encrypt the data
        byte[] encryptedData = cipher.doFinal(data);

        //Concatenate everything and return the final data
        ByteBuffer byteBuffer = ByteBuffer.allocate(4 + iv.length + encryptedData.length);
        byteBuffer.putInt(iv.length);
        byteBuffer.put(iv);
        byteBuffer.put(encryptedData);
        return byteBuffer.array();
    }


    public static byte[] decryptData(String key, byte[] encryptedData)
            throws NoSuchPaddingException,
            NoSuchAlgorithmException,
            InvalidAlgorithmParameterException,
            InvalidKeyException,
            BadPaddingException,
            IllegalBlockSizeException, InvalidKeySpecException {


        //Wrap the data into a byte buffer to ease the reading process
        ByteBuffer byteBuffer = ByteBuffer.wrap(encryptedData);

        int nonceSize = byteBuffer.getInt();

        //Make sure that the file was encrypted properly
        if (nonceSize < 12 || nonceSize >= 16) {
            throw new IllegalArgumentException("Nonce size is incorrect. Make sure that the incoming data is an AES encrypted file.");
        }
        byte[] iv = new byte[nonceSize];
        byteBuffer.get(iv);

        //Prepare your key/password
        SecretKey secretKey = generateSecretKey(key, iv);

        //get the rest of encrypted data
        byte[] cipherBytes = new byte[byteBuffer.remaining()];
        byteBuffer.get(cipherBytes);

        Cipher cipher = Cipher.getInstance("AES/GCM/NoPadding");
        GCMParameterSpec parameterSpec = new GCMParameterSpec(GCM_TAG_LENGTH * 8, iv);

        //Encryption mode on!
        cipher.init(Cipher.DECRYPT_MODE, secretKey, parameterSpec);

        //Encrypt the data
        return cipher.doFinal(cipherBytes);

    }

    /**
     * Function to generate a 256 bit key from the given password and iv
     *
     * @param password
     * @param iv
     * @return Secret key
     */
    public static SecretKey generateSecretKey(String password, byte[] iv) throws NoSuchAlgorithmException, InvalidKeySpecException {
        KeySpec spec = new PBEKeySpec(password.toCharArray(), iv, 65536, AES_KEY_SIZE);
        SecretKeyFactory secretKeyFactory = SecretKeyFactory.getInstance("PBKDF2WithHmacSHA1");
        byte[] key = secretKeyFactory.generateSecret(spec).getEncoded();
        return new SecretKeySpec(key, "AES");
    }


    public static byte[] readFile(String path) {
        File file = new File(path);
        byte[] fileData = new byte[(int) file.length()];
        try (FileInputStream fileInputStream = new FileInputStream(file)) {
            fileInputStream.read(fileData);
        } catch (IOException e) {
            e.printStackTrace();
        }
        return fileData;
    }

    public static void writeFile(String path, byte[] data) {
        try (FileOutputStream fileOutputStream = new FileOutputStream(path)) {
            fileOutputStream.write(data);
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    /**
     * 加密数据, 并导出 *.csv 文件
     *
     * @param data List<Map<String, Object>>
     * @param path 导出文件路径
     */
    public static void encryptDataToCsvFile(List<Map<String, Object>> data, String path) {
        Map<String, Object> stringObjectMap = data.get(0);
        Set<String> title = stringObjectMap.keySet();
        StringBuilder stringBuilder = new StringBuilder()
                .append(String.join(",", title))
                .append("\n");

        for (Map<String, Object> map : data) {
//            map.values().forEach(value -> stringBuffer.append(value).append(","));
            String mapResult = map.values().stream().map(String::valueOf).collect(joining(","));
            stringBuilder.append(mapResult).append("\n");
        }
        byte[] resultBytes = null;
        try {
            resultBytes = encryptData(KEY, stringBuilder.toString().getBytes(StandardCharsets.UTF_8));
        } catch (Exception e) {
            e.printStackTrace();
        }
        writeFile(path, resultBytes);
    }

    /**
     * 解密 *.csv 文件, 生成格式化数据
     *
     * @param sourcePath
     * @return Multimap<String, Object>
     */
    public static Multimap<String, Object> decryptCsvFileToObject(String sourcePath) {
        Multimap<String, Object> multimap = MultimapBuilder.hashKeys().arrayListValues().build();

        try {
            byte[] fileBytes = readFile(sourcePath);
            byte[] resultBytes = decryptData(KEY, fileBytes);
            InputStream targetStream = ByteSource.wrap(resultBytes).openStream();
            BufferedReader bufferedReader = new BufferedReader(new InputStreamReader(targetStream));
            Object[] objects = bufferedReader.lines().toArray();

            multimap.put("title", Arrays.stream(objects).limit(1).collect(toList()));
            Arrays.stream(objects)
                    .skip(1)
                    .map(obj ->
                            new ArrayList<Object>(Arrays.asList(obj.toString().split(","))))
                    .forEach(obj -> multimap.put("values", obj));
        } catch (Exception e) {
            e.printStackTrace();
        }
        return multimap;
    }



    /**
     * 加密文件
     *
     * @param source
     * @param target
     */
    public static void encryptFile(File source, File target) {
        doFile(source.getPath(), target.getPath(), 0);
    }

    /**
     * 解密文件
     *
     * @param source
     * @param target
     */
    public static void decryptFile(File source, File target) {
        doFile(source.getPath(), target.getPath(), 1);
    }

    private static void doFile(String sourcePath, String targetPath, Integer type) {
        try {
            byte[] fileBytes = readFile(sourcePath);
            byte[] resultBytes = null;
            if (type == 0) {
                resultBytes = encryptData(KEY, fileBytes);
            } else {
                resultBytes = decryptData(KEY, fileBytes);
            }
            writeFile(targetPath, resultBytes);
        } catch (Exception e) {
            e.printStackTrace();
        }
    }
}
